#author hejie@ramaxel.com
#version 1.0
#2025-05-30

import time
import paramiko
import subprocess
from typing import Dict, List, Optional, Tuple
import datetime
import re



def print_time(ctx: str):
    tstring = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f'{tstring}, {ctx}')
    


class BMCTester:
    

    def __init__(self, bmc_ip: str, bmc_username: str, bmc_password: str, 
                 os_username: str, os_password: str, os_ip: str):
        """
        初始化BMC测试器
        
        :param bmc_ip: BMC管理IP地址
        :param bmc_username: BMC用户名
        :param bmc_password: BMC密码
        :param os_username: 操作系统用户名
        :param os_password: 操作系统密码
        :param os_ip: 操作系统IP地址
        """
        self.bmc_ip = bmc_ip
        self.bmc_username = bmc_username
        self.bmc_password = bmc_password
        self.os_username = os_username
        self.os_password = os_password
        self.os_ip = os_ip
        self.ssh_client: Optional[paramiko.SSHClient] = None
        self.command_log_file = "os_command_logs.txt"  # 命令执行日志文件
        self.memory_log_file = "memory_check_logs.txt"  # 内存检查专用日志文件
        self.ipmitool_path = "ipmitool"


    def run_ipmitool(self, command: str) -> Tuple[bool, str]:
        """
        执行ipmitool命令
        
        :param command: ipmitool命令（不包括ipmitool本身）
        :return: (是否成功, 输出结果)
        """
        full_cmd = f"{self.ipmitool_path} -I lanplus -H {self.bmc_ip} -U {self.bmc_username} -P {self.bmc_password} {command}"
        try:
            result = subprocess.run(full_cmd, shell=True, check=True, 
                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                   text=True, timeout=30)
            return True, result.stdout.strip()
        except subprocess.CalledProcessError as e:
            return False, f"Error: {e.stderr.strip()}"
        except Exception as e:
            return False, f"Exception: {str(e)}"

    def get_power_status(self) -> str:
        """
        获取当前电源状态
        
        :return: 电源状态（on/off/unknown）
        """
        success, output = self.run_ipmitool("power status")
        if success:
            match = re.search(r"Chassis Power is (on|off)", output, re.IGNORECASE)
            if match:
                return match.group(1).lower()
        return "unknown"

    def send_power_command(self, command: str) -> bool:
        """
        发送电源控制命令
        
        :param command: power on/power off/power reset/power cycle
        :return: 是否成功
        """
        success, output = self.run_ipmitool(f"power {command}")
        if success:
            print_time(f"电源命令 '{command}' 执行成功")
            return True
        else:
            print_time(f"电源命令 '{command}' 执行失败: {output}")
            return False

    def connect_os(self) -> bool:
        """连接到操作系统SSH"""
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.ssh_client.connect(
                hostname=self.os_ip,
                username=self.os_username,
                password=self.os_password,
                timeout=10
            )
            return True
        except Exception as e:
            print(f"连接操作系统SSH失败: {str(e)}")
            return False

    def disconnect(self):
        """断开所有连接"""
        if self.ssh_client:
            self.ssh_client.close()
            self.ssh_client = None

    def wait_for_os(self, timeout: int = 300, interval: int = 10) -> bool:
        """
        等待操作系统启动并可用
        
        :param timeout: 超时时间(秒)
        :param interval: 检查间隔(秒)
        :return: 是否成功
        """
        print_time(f"等待操作系统启动，超时时间: {timeout}秒...")
        start_time = time.time()
        last_status = ""
        
        while time.time() - start_time < timeout:
            # 检查电源状态
            power_status = self.get_power_status()
            if power_status != "on":
                status = f"电源状态: {power_status}"
                if status != last_status:
                    print(status)
                    last_status = status
                time.sleep(interval)
                continue
            
            # 尝试SSH连接
            if not self.connect_os():
                status = "系统正在启动中，SSH不可用..."
                if status != last_status:
                    print(status)
                    last_status = status
                time.sleep(interval)
                continue
            
            # 检查系统负载和运行时间
            uptime_cmd = "cat /proc/uptime"
            stdin, stdout, stderr = self.ssh_client.exec_command(uptime_cmd)
            uptime_output = stdout.read().decode().strip()
            
            if uptime_output:
                uptime_seconds = float(uptime_output.split()[0])
                if uptime_seconds < 60:
                    status = f"系统已启动，运行时间: {uptime_seconds:.1f}秒，等待系统完全初始化..."
                    if status != last_status:
                        print_time(status)
                        last_status = status
                    time.sleep(interval)
                    continue
            
            # 检查关键系统服务
            services_cmd = "systemctl is-active multi-user.target"
            stdin, stdout, stderr = self.ssh_client.exec_command(services_cmd)
            services_output = stdout.read().decode().strip()
            
            if services_output == "active":
                print_time("系统已完全启动并可用")
                
                # 记录内存信息
                self.check_and_log_memory()
                return True
            
            status = "等待系统服务完全启动..."
            if status != last_status:
                print_time(status)
                last_status = status
            
            time.sleep(interval)
        
        print_time("等待操作系统超时")  
        return True #return False  #timeout 超时之后也要继续，往下走下一个cycle hejie

    def check_and_log_memory(self):
        """检查内存容量并记录到日志"""
        commands = [
            'free -h',
            'dmidecode -t memory',
            'cat /proc/meminfo',
            'ipmitool sel list'    #hejie add 
        ]
        results = self.execute_and_log_os_commands(commands, self.memory_log_file)
        
        if results:
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            with open(self.memory_log_file, "a") as f:
                f.write(f"\n=== 内存检查 ===\n")
                f.write(f"检查时间: {timestamp}\n")
                f.write(f"系统IP: {self.os_ip}\n")
                
                for cmd, output, error, status in results:
                    f.write(f"\n命令: {cmd}\n")
                    f.write(f"状态: {status}\n")
                    if output:
                        f.write(f"输出:\n{output}\n")
                    if error:
                        f.write(f"错误:\n{error}\n")
                
                f.write("="*50 + "\n")
            print_time("内存信息已记录到日志")

    def execute_and_log_os_commands(self, commands: List[str], log_file: Optional[str] = None) -> List[Tuple[str, str, str, str]]:
        """
        执行操作系统命令并记录返回值和执行时间
        
        :param commands: 要执行的命令列表
        :param log_file: 指定日志文件，如果为None则使用默认文件
        :return: 返回元组列表(命令, 标准输出, 标准错误, 执行状态)
        """
        if not self.ssh_client:
            if not self.connect_os():
                print_time("无法连接到操作系统SSH")
                return []
        
        log_file = log_file or self.command_log_file
        results = []
        
        for cmd in commands:
            try:
                start_time = time.time()
                stdin, stdout, stderr = self.ssh_client.exec_command(cmd)
                exit_status = stdout.channel.recv_exit_status()
                output = stdout.read().decode().strip()
                error = stderr.read().decode().strip()
                end_time = time.time()
                
                execution_time = f"{end_time - start_time:.2f}s"
                status = "Success" if exit_status == 0 else f"Failed (Exit code: {exit_status})"
                
                # 记录结果
                result = (cmd, output, error, status)
                results.append(result)
                
                # 写入日志文件
                timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                with open(log_file, "a") as f:
                    f.write(f"\n=== Command Execution Log ===\n")
                    f.write(f"Timestamp: {timestamp}\n")
                    f.write(f"Command: {cmd}\n")
                    f.write(f"Execution Time: {execution_time}\n")
                    f.write(f"Status: {status}\n")
                    if output:
                        f.write(f"Output:\n{output}\n")
                    if error:
                        f.write(f"Error:\n{error}\n")
                    f.write("="*40 + "\n")
                
                print_time(f"执行命令: {cmd} -> 状态: {status}, 耗时: {execution_time}")
                
                #add parse mce,uce,ecc hejie

            except Exception as e:
                error_msg = f"执行命令 '{cmd}' 时出错: {str(e)}"
                print_time(error_msg)
                results.append((cmd, "", error_msg, "Failed (Exception)"))
                
                # 记录错误到日志文件
                timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                with open(log_file, "a") as f:
                    f.write(f"\n=== Command Execution Log ===\n")
                    f.write(f"Timestamp: {timestamp}\n")
                    f.write(f"Command: {cmd}\n")
                    f.write(f"Status: Failed (Exception)\n")
                    f.write(f"Error: {error_msg}\n")
                    f.write("="*40 + "\n")


        
        return results

    def run_test_cycle(self, test_config: Dict) -> bool:
        """
        运行测试循环（根据配置执行测试项目）
        
        :param test_config: 测试配置字典
        :return: 是否全部成功
        """
        test_results = []
        cycles = test_config.get("cycles", 1)
        delay = test_config.get("delay_between_cycles", 30)
        timeTemp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')


        for i in range(1, cycles + 1):

            time1 = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            timestamp1 = time.time()

            print_time(f" === 开始测试循环 {i}/{cycles} ===")
            
            # 打印当前电源状态
            print_time("[1/3] 测试开机...")

            power_status = self.get_power_status()
            print_time(f"当前电源状态: {power_status}")
            
            # 测试开机（如果配置启用）
            if test_config.get("test_power_on", True):
			    
                if not self.send_power_command("on"):
                    test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}, 循环 {i}: 开机失败")
                    continue
                
                if not self.wait_for_os():
                    test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}, 循环 {i}: 开机后系统未就绪")
                    continue
                
                # 系统启动后执行命令（如果配置启用）
                if test_config.get("execute_os_commands", False):
                    commands = test_config.get("os_commands", [
                        'hostname',
                        'uname -a',
                        'uptime',
                        'df -h',
                        'free -m',
                        'ipmitool sel clear'  #hejie
                    ])
                    print_time(f"\n 执行操作系统命令并记录结果...")
                    self.execute_and_log_os_commands(commands)
            
            # 测试重启/DC（如果配置启用）
            test_type = test_config.get("test_type", "reboot")  # reboot或dc
            if test_config.get("test_reboot", True):
                print(f"\n")
                #timeTemp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                print_time(f"[2/3] 测试{test_type.upper()}...")
                if test_type == "dc":
                    # DC测试：先关机再开机
                    if not self.send_power_command("off"):
                        test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}: DC测试关机失败")
                        continue
                    time.sleep(60)  #time.sleep(10)  # 等待完全关机
                    if not self.send_power_command("on"):
                        test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}: DC测试开机失败")
                        continue
                else:
                    # 普通重启
                    if not self.send_power_command("reset"):
                        test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}: 重启失败")
                        continue
                
                # if not self.wait_for_os():
                #     time2 = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                #     timestamp2 = time.time()
                    
                #     if( time2 - time1 > 200):  # if timeout must goto next cycle , hejie 2025-05-29
                #         test_results.append(f"循环 {i}/{cycles}: {test_type.upper()}后系统未就绪,超时200s,准备下一次DC {time2}")
                #         break
                #     else:
                #         test_results.append(f"循环 {i}/{cycles}: {test_type.upper()}后系统未就绪 {time2}")
                #         continue

                if not self.wait_for_os():
                    test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}/{cycles}: {test_type.upper()}后系统未就绪")
                    continue
                    
                # 系统重启后执行命令（如果配置启用）
                if test_config.get("execute_os_commands", False):
                    commands = test_config.get("os_commands", [
                        'hostname',
                        'uname -a',
                        'uptime',
                        'dmesg | tail -n 20',
                        'cat /proc/uptime',
                        'ipmitool sel list',  #hejie
                        'ipmitool sel clear'
                        #'dmesg',
                        #'dmesg -c'
                    ])
                    print_time("执行操作系统命令并记录结果...")
                    self.execute_and_log_os_commands(commands)
                
                # 特别记录重启后的内存信息
                self.check_and_log_memory()
            
            # 测试关机（如果配置启用）
            if test_config.get("test_power_off", True):
                print_time("[3/3] 测试关机...")
                if not self.send_power_command("off"):
                    test_results.append(f"循环 {i}: 关机失败")
                    continue
                
                # 等待关机完成 30,60
                time.sleep(30)#time.sleep(30)
            
            # 准备下一次循环（如果不是最后一个循环）
            if i < cycles:
                print_time("准备下一次测试循环...")
                if test_config.get("test_power_on", True):  # 只有启用了开机测试才需要准备
                    if not self.send_power_command("on"):
                        test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}: 准备下次循环时开机失败")
                        break
                    
                    if not self.wait_for_os():
                        test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}: 准备下次循环时系统未就绪")
                        break
                    
                    # 执行命令（如果配置启用）
                    if test_config.get("execute_os_commands", False):
                        commands = test_config.get("os_commands", [
                            'hostname',
                            'date',
                            'cat /etc/os-release',
                            'ip a',
                            'ipmitool sel list',
                            'ipmitool sel clear'
                        ])
                        print_time("\n执行操作系统命令并记录结果...")
                        self.execute_and_log_os_commands(commands)
                    
                    # 记录内存信息
                    self.check_and_log_memory()
            
            test_results.append(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')},循环 {i}: 成功完成")
            time.sleep(delay)
        
        # 打印结果
        print_time("\n=== 测试结果 ===")
        for result in test_results:
            print_time(f"- {result}")
        
        self.disconnect()
        return all("成功" in result for result in test_results)

if __name__ == "__main__":
    # 配置参数（建议从环境变量或配置文件读取）
    config = {
        # 连接配置
        "BMC_IP": "192.168.67.59",
        "BMC_USERNAME": "taobao",
        "BMC_PASSWORD": "9ijn0okm",
        "OS_IP": "192.168.67.123",
        "OS_USERNAME": "root",
        "OS_PASSWORD": "1234",
        
        # 测试配置
        "test_options": {
            "cycles": 100,                      # 测试循环次数
            "delay_between_cycles": 30,       # 循环间隔(秒)
            "test_power_on": False,           # 是否测试开机
            "test_power_off": False,          # 是否测试关机
            "test_reboot": True,             # 是否测试重启/DC
            "test_type": "dc",               # 测试类型: reboot或dc
            "execute_os_commands": True,     # 是否执行OS命令
            "os_commands": [                # 自定义要执行的命令列表
                'hostname',
                'uname -a',
                'uptime',
                'df -h',
                'free -m',
                'lsblk',
                'ip a'
            ]
        }
    }
    
    # 创建测试器实例
    tester = BMCTester(
        bmc_ip=config["BMC_IP"],
        bmc_username=config["BMC_USERNAME"],
        bmc_password=config["BMC_PASSWORD"],
        os_ip=config["OS_IP"],
        os_username=config["OS_USERNAME"],
        os_password=config["OS_PASSWORD"]
    )
    
    # 运行测试（传入测试配置）
    success = tester.run_test_cycle(test_config=config["test_options"])
    
    if success:
        print_time("\n✅ 所有测试循环成功完成!")
        print_time(f"操作系统命令执行日志已保存到: {tester.command_log_file}")
        print_time(f"内存检查日志已保存到: {tester.memory_log_file}")
    else:
        print_time("\n❌ 测试过程中出现故障，请检查日志!")